﻿#pragma once
#include <iostream>
#include <assert.h>
using namespace std;
template<class T>
struct AVLTreeNode
{
	AVLTreeNode(const T& data = T())
		: _pLeft(nullptr)
		, _pRight(nullptr)
		, _pParent(nullptr)
		, _data(data)
		, _bf(0)
	{}

	AVLTreeNode<T>* _pLeft;
	AVLTreeNode<T>* _pRight;
	AVLTreeNode<T>* _pParent;
	T _data;
	int _bf;   // 节点的平衡因子
};


// AVL: 二叉搜索树 + 平衡因子的限制
template<class T>
class AVLTree
{
	typedef AVLTreeNode<T> Node;
public:

	AVLTree()
		: _pRoot(nullptr)
	{}

	// 在AVL树中插入值为data的节点
	bool Insert(const T& kv)
	{
		if (_pRoot == nullptr)
		{
			_pRoot = new Node(kv);
			return true;
		}
		Node* parent = nullptr;
		Node* cur = _pRoot;
		while (cur)
		{
			if (cur->_data < kv)
			{
				parent = cur;
				cur = cur->_pRight;
			}
			else if (cur->_data > kv)
			{
				parent = cur;
				cur = cur->_pLeft;
			}
			else
			{
				return false;
			}
		}
		cur = new Node(kv);
		if (parent->_data < kv)
		{
			parent->_pRight = cur;
		}
		else
		{
			parent->_pLeft = cur;
		}
		// 链接父亲
		cur->_pParent = parent;
		// 控制平衡
		// 更新平衡因子
		while (parent)
		{
			if (cur == parent->_pLeft)
				parent->_bf--;
			else
				parent->_bf++;

			if (parent->_bf == 0)
			{
				break;
			}
			else if (parent->_bf == 1 || parent->_bf == -1)
			{
				cur = parent;
				parent = parent->_pParent;
			}
			else if (parent->_bf == 2 || parent->_bf == -2)
			{
				if (parent->_bf == -2 && cur->_bf == -1)
				{
					RotateR(parent);
				}
				else if (parent->_bf == 2 && cur->_bf == 1)
				{
					RotateL(parent);
				}
				else if (parent->_bf == -2 && cur->_bf == 1)
				{
					RotateLR(parent);
				}
				else if (parent->_bf == 2 && cur->_bf == -1)
				{
					RotateRL(parent);
				}
				else
				{
					assert(false);
				}
				break;
			}
			else
			{
				assert(false);
			}
		}
		return true;

	}
		// AVL树的验证
		bool IsAVLTree()
		{
			return _IsAVLTree(_pRoot);
		}

		void InOrder()
		{
			inOrderTraversal(_pRoot);
			cout << endl;
		}

		int Height()
		{
			return _Height(_pRoot);
		}
		void inOrderTraversal(Node * root)
		{
			if (root != nullptr) {
				inOrderTraversal(root->_pLeft);
				cout << root->_data << " ";
				inOrderTraversal(root->_pRight);
			}
		}

		Node* Find(const T & key)
		{
			Node* cur = _pRoot;
			while (cur)
			{
				if (cur->_data < key)
				{
					cur = cur->_pRight;
				}
				else if (cur->_data > key)
				{
					cur = cur->_pLeft;
				}
				else
				{
					return cur;
				}
			}
			return nullptr;
		}

	private:
		// 根据AVL树的概念验证pRoot是否为有效的AVL树
		//左右节点因子数
		bool _IsAVLTree(Node * pRoot)
		{
			if (pRoot == nullptr)
				return true;
			int L = _Height(pRoot->_pLeft);
			int R = _Height(pRoot->_pRight);
			if (L - R > 1 || L - R < -1)
			{
				cout << pRoot->_data << "高度差异常" << endl;
				return false;
			}
			else if ((R - L) != pRoot->_bf)
			{
				cout << pRoot->_data << "平衡因子异常" << pRoot->_bf << endl;
				return false;
			}
			return _IsAVLTree(pRoot->_pLeft) && _IsAVLTree(pRoot->_pRight);
		}

		size_t _Height(Node * pRoot)
		{
			if (pRoot == nullptr)
			{
				return 0;
			}
			int L = _Height(pRoot->_pLeft);
			int H = _Height(pRoot->_pRight);
			return L > H ? L + 1 : H + 1;
		}

		// 右单旋
		void RotateR(Node * pParent)
		{
			Node* SubL = pParent->_pLeft;
			Node* SubLR = SubL->_pRight;

			pParent->_pLeft = SubLR;
			if (SubLR)
			{
				SubLR->_pParent = pParent;
			}
			SubL->_pRight = pParent;
			Node* Gparent = pParent->_pParent;
			pParent->_pParent = SubL;

			if (Gparent == nullptr)
			{
				_pRoot = SubL;
			}
			else
			{
				if (Gparent->_pLeft == pParent)
				{
					Gparent->_pLeft = SubL;
				}
				else
				{
					Gparent->_pRight = SubL;
				}
			}
			SubL->_pParent = Gparent;
			pParent->_bf = SubL->_bf = 0;
		}

		// 左单旋
		void RotateL(Node * pParent)
		{
			Node* SubR = pParent->_pRight;
			Node* SubRL = SubR->_pLeft;
			pParent->_pRight = SubRL;
			if (SubRL)
			{
				SubRL->_pParent = pParent;
			}
			SubR->_pLeft = pParent;
			Node* Gparent = pParent->_pParent;
			pParent->_pParent = SubR;

			if (Gparent == nullptr)
			{
				_pRoot = SubR;
			}
			else
			{
				if (Gparent->_pLeft == pParent)
				{
					Gparent->_pLeft = SubR;
				}
				else
				{
					Gparent->_pRight = SubR;
				}
			}
			SubR->_pParent = Gparent;
			pParent->_bf = SubR->_bf = 0;
		}

		// 右左双旋
		void RotateRL(Node * pParent)
		{
			Node* SubR = pParent->_pRight;
			Node* SubRL = SubR->_pLeft;
			int bf = SubRL->_bf;
			RotateR(SubR);
			RotateL(pParent);
			//

			//SubR->_pLeft = SubRL->_pRight;
			//if (SubRL->_pRight)
			//	SubR->_pRight->_pParent = SubR;
			//SubRL->_pRight = SubR;
			//SubR->_pParent = SubRL;

			//pParent->_pRight = SubRL->_pLeft;
			//if (SubRL->_pLeft)
			//	SubRL->_pLeft->_pParent = pParent;
			//SubRL->_pLeft = pParent;
			//SubRL->_pParent = pParent->_pParent
			if (bf == 0)
			{
				SubR->_bf = 0;
				SubRL->_bf = 0;
				pParent->_bf = 0;
			}
			else if (bf == -1)
			{
				SubR->_bf = 1;
				SubRL->_bf = 0;
				pParent->_bf = 0;
			}
			else if (bf == 1)
			{
				SubR->_bf = 0;
				SubRL->_bf = 0;
				pParent->_bf = -1;
			}
			else
			{
				assert(false);
			}
		}
	

	// 左右双旋
	void RotateLR(Node* pParent)
	{
		Node* SubL = pParent->_pLeft;
		Node* SubLR = SubL->_pRight;
		int bf = SubLR->_bf;

		RotateL(SubL);
		RotateR(pParent);
		if (bf == 0)
		{
			SubL->_bf = 0;
			SubLR->_bf = 0;
			pParent->_bf = 0;
		}
		else if (bf == -1)
		{
			SubL->_bf = 0;
			SubLR->_bf = 0;
			pParent->_bf = 1;
		}
		else if (bf == 1)
		{
			SubL->_bf = -1;
			SubLR->_bf = 0;
			pParent->_bf = 0;
		}
		else
		{
			assert(false);
		}
	}


private:
	Node* _pRoot;
};

